{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "201501f9",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.append('./')\n",
    "sys.path.append('../')\n",
    "sys.path.append('../..')\n",
    "\n",
    "import os\n",
    "import pandas as pd\n",
    "from sklearn import preprocessing\n",
    "import string\n",
    "from typing import Sequence, Tuple, List, Union\n",
    "from tqdm import tqdm\n",
    "import fm\n",
    "import torch\n",
    "from torch import nn\n",
    "from torch import optim\n",
    "from torch.utils.data import DataLoader\n",
    "from torch.utils.data.sampler import RandomSampler\n",
    "import numpy as np\n",
    "import random\n",
    "\n",
    "def seed_torch(seed=0):\n",
    "    random.seed(seed)\n",
    "    os.environ['PYTHONHASHSEED'] = str(seed)\n",
    "    np.random.seed(seed)\n",
    "    torch.manual_seed(seed)\n",
    "    torch.cuda.manual_seed(seed)\n",
    "    torch.cuda.manual_seed_all(seed) # if you are using multi-GPU.\n",
    "    torch.backends.cudnn.benchmark = False\n",
    "    torch.backends.cudnn.deterministic = True\n",
    "\n",
    "seed_torch(2021)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "36614033",
   "metadata": {},
   "source": [
    "## 1. Load Model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "af3952d7",
   "metadata": {},
   "source": [
    "### (1) define utr_function_predictor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "b5aee6b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Human5PrimeUTRPredictor(torch.nn.Module):\n",
    "    \"\"\"\n",
    "    contact predictor with inner product\n",
    "    \"\"\"\n",
    "    def __init__(self, task=\"rgs\", arch=\"cnn\", input_types=[\"seq\", \"emb-rnafm\"]):\n",
    "        \"\"\"\n",
    "        :param depth_reduction: mean, first\n",
    "        \"\"\"       \n",
    "        super().__init__()        \n",
    "        self.task = task\n",
    "        self.arch = arch\n",
    "        self.input_types = input_types        \n",
    "        self.padding_mode = \"right\" #\"left\"\n",
    "        self.token_len = 100\n",
    "        self.out_plane = 1\n",
    "        self.in_channels = 0\n",
    "        if \"seq\" in self.input_types:\n",
    "            self.in_channels = self.in_channels + 4\n",
    "\n",
    "        if \"emb-rnafm\" in self.input_types:\n",
    "            self.reductio_module = nn.Linear(640, 32)\n",
    "            self.in_channels = self.in_channels + 32  \n",
    "\n",
    "        if self.arch == \"cnn\":\n",
    "            self.predictor = self.create_1dcnn_for_emd(in_planes=self.in_channels, out_planes=1)\n",
    "\n",
    "    def forward(self, tokens, inputs):\n",
    "        ensemble_inputs = []\n",
    "        if \"seq\" in inputs:\n",
    "            nest_token = (tokens[:, 1:-1] - 4)\n",
    "\n",
    "            # 100 right padding\n",
    "            nest_token = torch.nn.functional.pad(nest_token, (0, self.token_len - nest_token.shape[1]), value=-1)\n",
    "            token_padding_mask = nest_token.ge(0).long()\n",
    "            one_hot_tokens = torch.nn.functional.one_hot((nest_token * token_padding_mask),\n",
    "                                                         num_classes=4).float() * token_padding_mask.unsqueeze(-1)\n",
    "            if self.arch == \"cnn\":     # B, L, 4\n",
    "                one_hot_tokens = one_hot_tokens.permute(0, 2, 1)\n",
    "            else:\n",
    "                one_hot_tokens = one_hot_tokens.view(one_hot_tokens.shape[0], -1)\n",
    "            ensemble_inputs.append(one_hot_tokens)\n",
    "\n",
    "        if \"emb-rnafm\" in inputs:\n",
    "            embeddings = inputs[\"emb-rnafm\"]\n",
    "            # remove auxiliary tokens\n",
    "            embeddings, padding_masks = self.remove_pend_tokens_1d(tokens, embeddings)\n",
    "            if len(embeddings.size()) == 3:  # for pure seq\n",
    "                batch_size, seqlen, hiddendim = embeddings.size()\n",
    "            elif len(embeddings.size()) == 4:  # for msa\n",
    "                batch_size, depth, seqlen, hiddendim = embeddings.size()\n",
    "                # reduction\n",
    "                embeddings = self.msa_depth_reduction(embeddings, padding_masks)\n",
    "            else:\n",
    "                raise Exception(\"Unknown Embedding Type!\")\n",
    "\n",
    "            # 100 right padding\n",
    "            embeddings = torch.nn.functional.pad(embeddings, (0, 0, 0, self.token_len - embeddings.shape[1]))\n",
    "            embeddings = self.reductio_module(embeddings)\n",
    "\n",
    "            if self.arch == \"cnn\":\n",
    "                embeddings = embeddings.permute(0, 2, 1)\n",
    "            else:\n",
    "                embeddings = embeddings.reshape(embeddings.shape[0], -1)\n",
    "\n",
    "            ensemble_inputs.append(embeddings)        \n",
    "\n",
    "        if len(ensemble_inputs) > 0:\n",
    "            ensemble_inputs = torch.cat(ensemble_inputs, dim=1)\n",
    "\n",
    "        if self.padding_mode == \"left\":\n",
    "            # 100 left padding\n",
    "            nest_token = (tokens[:, 1:-1] - 4)\n",
    "            nest_token = torch.nn.functional.pad(nest_token, (0, self.token_len - nest_token.shape[1]), value=-1)\n",
    "            token_padding_mask = nest_token.ge(0).long()\n",
    "\n",
    "            left_ensembles = []\n",
    "            for i in range(nest_token.shape[0]):\n",
    "                length = token_padding_mask[i].sum().item()\n",
    "                left_ensemble = torch.cat([ensemble_inputs[i, :, length:], ensemble_inputs[i, :, 0: length]], dim=-1)\n",
    "                left_ensembles.append(left_ensemble)\n",
    "            ensemble_inputs = torch.stack(left_ensembles, dim=0)\n",
    "\n",
    "        if isinstance(ensemble_inputs, list) != True:\n",
    "            output = self.predictor(ensemble_inputs)\n",
    "            output = output.squeeze(-1)\n",
    "        else:\n",
    "            output = 0\n",
    "        return output\n",
    " \n",
    "    def create_1dcnn_for_emd(self, in_planes, out_planes):\n",
    "        main_planes = 64\n",
    "        dropout = 0.2\n",
    "        emb_cnn = nn.Sequential(\n",
    "            nn.Conv1d(in_planes, main_planes, kernel_size=3, padding=1),  ## 3\n",
    "            ResBlock(main_planes * 1, main_planes * 1, stride=2, dilation=1, conv_layer=nn.Conv1d,\n",
    "                     norm_layer=nn.BatchNorm1d), \n",
    "            ResBlock(main_planes * 1, main_planes * 1, stride=1, dilation=1, conv_layer=nn.Conv1d,\n",
    "                     norm_layer=nn.BatchNorm1d),  \n",
    "            ResBlock(main_planes * 1, main_planes * 1, stride=2, dilation=1, conv_layer=nn.Conv1d,\n",
    "                     norm_layer=nn.BatchNorm1d), \n",
    "            ResBlock(main_planes * 1, main_planes * 1, stride=1, dilation=1, conv_layer=nn.Conv1d,\n",
    "                     norm_layer=nn.BatchNorm1d),  \n",
    "            ResBlock(main_planes * 1, main_planes * 1, stride=2, dilation=1, conv_layer=nn.Conv1d,\n",
    "                     norm_layer=nn.BatchNorm1d), \n",
    "            ResBlock(main_planes * 1, main_planes * 1, stride=1, dilation=1, conv_layer=nn.Conv1d,\n",
    "                     norm_layer=nn.BatchNorm1d),       \n",
    "            nn.AdaptiveAvgPool1d(1),\n",
    "            nn.Flatten(),\n",
    "            nn.Dropout(dropout),\n",
    "            nn.Linear(main_planes * 1, out_planes),\n",
    "        )\n",
    "        return emb_cnn\n",
    "    \n",
    "    def remove_pend_tokens_1d(self, tokens, seqs):\n",
    "        \"\"\"\n",
    "        :param tokens:\n",
    "        :param seqs: must be shape of [B, ..., L, E]    # seq: [B, L, E]; msa: [B, D, L, E]\n",
    "        :return:\n",
    "        \"\"\"\n",
    "        self.pad_idx = 1\n",
    "        self.eos_idx = 2\n",
    "        self.append_eos = True\n",
    "        self.prepend_bos = True\n",
    "        \n",
    "        padding_masks = tokens.ne(self.pad_idx)\n",
    "\n",
    "        # remove eos token  （suffix first）\n",
    "        if self.append_eos:     # default is right\n",
    "            eos_masks = tokens.ne(self.eos_idx)\n",
    "            eos_pad_masks = (eos_masks & padding_masks).to(seqs)\n",
    "            seqs = seqs * eos_pad_masks.unsqueeze(-1)\n",
    "            seqs = seqs[:, ..., :-1, :]\n",
    "            padding_masks = padding_masks[:, ..., :-1]\n",
    "\n",
    "        # remove bos token\n",
    "        if self.prepend_bos:    # default is left\n",
    "            seqs = seqs[:, ..., 1:, :]\n",
    "            padding_masks = padding_masks[:, ..., 1:]\n",
    "\n",
    "        if not padding_masks.any():\n",
    "            padding_masks = None\n",
    "\n",
    "        return seqs, padding_masks\n",
    "\n",
    "class ResBlock(nn.Module):\n",
    "    def __init__(\n",
    "        self,\n",
    "        in_planes,\n",
    "        out_planes,\n",
    "        stride=1,\n",
    "        dilation=1,\n",
    "        conv_layer=nn.Conv2d,\n",
    "        norm_layer=nn.BatchNorm2d,\n",
    "    ):\n",
    "        super(ResBlock, self).__init__()        \n",
    "        self.bn1 = norm_layer(in_planes)\n",
    "        self.relu1 = nn.ReLU(inplace=True)\n",
    "        self.conv1 = conv_layer(in_planes, out_planes, kernel_size=3, stride=stride, padding=dilation, bias=False)       \n",
    "        self.bn2 = norm_layer(out_planes)\n",
    "        self.relu2 = nn.ReLU(inplace=True)\n",
    "        self.conv2 = conv_layer(out_planes, out_planes, kernel_size=3, padding=dilation, bias=False)\n",
    "\n",
    "        if stride > 1 or out_planes != in_planes: \n",
    "            self.downsample = nn.Sequential(\n",
    "                conv_layer(in_planes, out_planes, kernel_size=1, stride=stride, bias=False),\n",
    "                norm_layer(out_planes),\n",
    "            )\n",
    "        else:\n",
    "            self.downsample = None\n",
    "\n",
    "        self.stride = stride\n",
    "\n",
    "    def forward(self, x):\n",
    "        identity = x\n",
    "\n",
    "        out = self.bn1(x)\n",
    "        out = self.relu1(out)\n",
    "        out = self.conv1(out)        \n",
    "        out = self.bn2(out)\n",
    "        out = self.relu2(out)\n",
    "        out = self.conv2(out)\n",
    "\n",
    "        if self.downsample is not None:\n",
    "            identity = self.downsample(x)\n",
    "\n",
    "        out += identity\n",
    "\n",
    "        return out\n",
    "    \n",
    "def weights_init(m):\n",
    "    classname = m.__class__.__name__\n",
    "    if classname.find('Linear') != -1:\n",
    "        nn.init.normal_(m.weight, std=0.001) #std=0.001)\n",
    "        if isinstance(m.bias, nn.Parameter):\n",
    "            nn.init.constant_(m.bias, 0.0)\n",
    "    elif classname.find('BasicConv') != -1: \n",
    "        pass\n",
    "    elif classname.find('Conv') != -1:\n",
    "        nn.init.kaiming_normal_(m.weight, a=0, mode='fan_in')\n",
    "        if m.bias is not None:\n",
    "            nn.init.constant_(m.bias, 0.0)\n",
    "    elif classname.find('BatchNorm') != -1:\n",
    "        if m.affine:\n",
    "            nn.init.constant_(m.weight, 1.0)\n",
    "            nn.init.constant_(m.bias, 0.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d28a72d6",
   "metadata": {},
   "source": [
    "### (2) create RNA-FM backbone"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "e1551e78",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "create RNA-FM_backbone sucessfully\n"
     ]
    }
   ],
   "source": [
    "device=\"cuda\"   # \"cpu\"\n",
    "os.environ['CUDA_VISIBLE_DEVICES'] = \"1\"  \n",
    "backbone, alphabet = fm.pretrained.esm1b_rna_t12(\"../../redevelop/pretrained/RNA-FM_pretrained.pth\")\n",
    "backbone.to(device)\n",
    "print(\"create RNA-FM_backbone sucessfully\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5d07fc1e",
   "metadata": {},
   "source": [
    "### (3) create UTR function downstream predictor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "97442953",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "create utr_func_predictor sucessfully\n",
      "Human5PrimeUTRPredictor(\n",
      "  (reductio_module): Linear(in_features=640, out_features=32, bias=True)\n",
      "  (predictor): Sequential(\n",
      "    (0): Conv1d(36, 64, kernel_size=(3,), stride=(1,), padding=(1,))\n",
      "    (1): ResBlock(\n",
      "      (bn1): BatchNorm1d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu1): ReLU(inplace=True)\n",
      "      (conv1): Conv1d(64, 64, kernel_size=(3,), stride=(2,), padding=(1,), bias=False)\n",
      "      (bn2): BatchNorm1d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu2): ReLU(inplace=True)\n",
      "      (conv2): Conv1d(64, 64, kernel_size=(3,), stride=(1,), padding=(1,), bias=False)\n",
      "      (downsample): Sequential(\n",
      "        (0): Conv1d(64, 64, kernel_size=(1,), stride=(2,), bias=False)\n",
      "        (1): BatchNorm1d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      )\n",
      "    )\n",
      "    (2): ResBlock(\n",
      "      (bn1): BatchNorm1d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu1): ReLU(inplace=True)\n",
      "      (conv1): Conv1d(64, 64, kernel_size=(3,), stride=(1,), padding=(1,), bias=False)\n",
      "      (bn2): BatchNorm1d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu2): ReLU(inplace=True)\n",
      "      (conv2): Conv1d(64, 64, kernel_size=(3,), stride=(1,), padding=(1,), bias=False)\n",
      "    )\n",
      "    (3): ResBlock(\n",
      "      (bn1): BatchNorm1d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu1): ReLU(inplace=True)\n",
      "      (conv1): Conv1d(64, 64, kernel_size=(3,), stride=(2,), padding=(1,), bias=False)\n",
      "      (bn2): BatchNorm1d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu2): ReLU(inplace=True)\n",
      "      (conv2): Conv1d(64, 64, kernel_size=(3,), stride=(1,), padding=(1,), bias=False)\n",
      "      (downsample): Sequential(\n",
      "        (0): Conv1d(64, 64, kernel_size=(1,), stride=(2,), bias=False)\n",
      "        (1): BatchNorm1d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      )\n",
      "    )\n",
      "    (4): ResBlock(\n",
      "      (bn1): BatchNorm1d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu1): ReLU(inplace=True)\n",
      "      (conv1): Conv1d(64, 64, kernel_size=(3,), stride=(1,), padding=(1,), bias=False)\n",
      "      (bn2): BatchNorm1d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu2): ReLU(inplace=True)\n",
      "      (conv2): Conv1d(64, 64, kernel_size=(3,), stride=(1,), padding=(1,), bias=False)\n",
      "    )\n",
      "    (5): ResBlock(\n",
      "      (bn1): BatchNorm1d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu1): ReLU(inplace=True)\n",
      "      (conv1): Conv1d(64, 64, kernel_size=(3,), stride=(2,), padding=(1,), bias=False)\n",
      "      (bn2): BatchNorm1d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu2): ReLU(inplace=True)\n",
      "      (conv2): Conv1d(64, 64, kernel_size=(3,), stride=(1,), padding=(1,), bias=False)\n",
      "      (downsample): Sequential(\n",
      "        (0): Conv1d(64, 64, kernel_size=(1,), stride=(2,), bias=False)\n",
      "        (1): BatchNorm1d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      )\n",
      "    )\n",
      "    (6): ResBlock(\n",
      "      (bn1): BatchNorm1d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu1): ReLU(inplace=True)\n",
      "      (conv1): Conv1d(64, 64, kernel_size=(3,), stride=(1,), padding=(1,), bias=False)\n",
      "      (bn2): BatchNorm1d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu2): ReLU(inplace=True)\n",
      "      (conv2): Conv1d(64, 64, kernel_size=(3,), stride=(1,), padding=(1,), bias=False)\n",
      "    )\n",
      "    (7): AdaptiveAvgPool1d(output_size=1)\n",
      "    (8): Flatten(start_dim=1, end_dim=-1)\n",
      "    (9): Dropout(p=0.2, inplace=False)\n",
      "    (10): Linear(in_features=64, out_features=1, bias=True)\n",
      "  )\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "task=\"rgs\"\n",
    "arch=\"cnn\"\n",
    "input_items = [\"seq\",\"emb-rnafm\"]   # [\"seq\"], [\"emb-rnafm\"], [\"seq\",\"emb-rnafm\"] \n",
    "model_name = arch.upper() + \"_\" + \"_\".join(input_items) \n",
    "utr_func_predictor = Human5PrimeUTRPredictor(\n",
    "    task=task, arch=arch, input_types=input_items    \n",
    ")\n",
    "utr_func_predictor.apply(weights_init)\n",
    "utr_func_predictor.to(device)\n",
    "print(\"create utr_func_predictor sucessfully\")\n",
    "print(utr_func_predictor)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "be5e47f0",
   "metadata": {},
   "source": [
    "### (4) define loss function and optimizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "8d1208bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "criterion = nn.MSELoss(reduction=\"none\")\n",
    "optimizer = optim.Adam(utr_func_predictor.parameters(), lr=0.001)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9d2c6444",
   "metadata": {},
   "source": [
    "## 2. Load Data\n",
    "You should download data from gdrive link: https://drive.google.com/file/d/10zCfOHOaOa__J2AIuZyidZ9sVJ9L11rI/view?usp=sharing and place it in the tutorials/utr-function-prediction"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2d679217",
   "metadata": {},
   "source": [
    "### (1) define utr dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "6be0f881",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Human_5Prime_UTR_VarLength(object):\n",
    "    def __init__(self, root, set_name=\"train\", input_items=[\"seq\", \"emb-rnafm\"], label_items=[\"rl\"], use_cache=0):\n",
    "        \"\"\"\n",
    "        :param root: root path of dataset - CATH. however not all of stuffs under this root path\n",
    "        :param data_type: seq, msa\n",
    "        :param label_type: 1d, 2d\n",
    "        :param set_name: \"train\", \"valid\", \"test\"\n",
    "        \"\"\"\n",
    "        self.root = root\n",
    "        self.theme = \"rna\"\n",
    "        self.data_type = \"seq\"\n",
    "        self.set_name = set_name          \n",
    "        self.use_cache = 1\n",
    "        self.cache = {}\n",
    "\n",
    "        self.input_items = input_items\n",
    "        if \"seq\" not in self.input_items:\n",
    "            self.input_items.append(\"seq\")  # for generate RNA-FM embedding\n",
    "        self.label_items = label_items\n",
    "        self.data_items = input_items + label_items  # combine two items\n",
    "    \n",
    "        # 1. Create Paths\n",
    "        self.src_scv_path = os.path.join(self.root, \"data\", \"GSM4084997_varying_length_25to100.csv\") \n",
    "        \n",
    "        # 2. Create Data INFO        \n",
    "        self.data_avatars, self.stats = self.__dataset_info(self.src_scv_path, self.data_items)\n",
    "        self.labels = ['rl']\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        inputs = {}\n",
    "        inputs[\"seq\"] = (self.data_avatars[\"name\"][index], self.data_avatars[\"seq\"][index][0].replace(\"T\", \"U\"))\n",
    "        labels = {\"rl\": self.data_avatars[\"rl\"][index],}\n",
    "        return inputs, labels\n",
    "\n",
    "    def __len__(self):\n",
    "        return self.stats.shape[0]\n",
    "\n",
    "    def __dataset_info(self, src_csv_path, data_items):\n",
    "        \"\"\"\n",
    "        :param name_path: txt record name list for specific set_name\n",
    "        :param data_dir:\n",
    "        :param msa_dir:\n",
    "        :param ann1d_dir:\n",
    "        :param ann2d_dir:\n",
    "        :return:\n",
    "        \"\"\"\n",
    "        seq_col = ['utr100']\n",
    "        label_col = ['rl']  # need scale\n",
    "\n",
    "        src_df = pd.read_csv(src_csv_path)   #, index_col=0)\n",
    "        src_df.loc[:, \"ori_index\"] = src_df.index\n",
    "        random_df = src_df[src_df['set'] == 'random']\n",
    "        ## Filter out UTRs with too few less reads\n",
    "        random_df = random_df[random_df['total_reads'] >= 10]    # 87000 -> 83919\n",
    "        random_df['utr100'] = random_df['utr'] # + 75 * 'N'  # left padding from reference code\n",
    "        #random_df['utr100'] = random_df['utr100'].str[-100:]\n",
    "        random_df.sort_values('total_reads', inplace=True, ascending=False)\n",
    "        random_df.reset_index(inplace=True, drop=True)\n",
    "\n",
    "        # human set for evaluation too\n",
    "        human_df = src_df[src_df['set'] == 'human']\n",
    "        ## Filter out UTRs with too few less reads\n",
    "        human_df = human_df[human_df['total_reads'] >= 10]   # 16739 -> 15555\n",
    "        human_df['utr100'] = human_df['utr'] #+ 75 * 'N'  # left padding from reference code, here we use right padding\n",
    "        #human_df['utr100'] = human_df['utr100'].str[-100:]\n",
    "        human_df.sort_values('total_reads', inplace=True, ascending=False)\n",
    "        human_df.reset_index(inplace=True, drop=True)\n",
    "\n",
    "        random_df_test = pd.DataFrame(columns=random_df.columns)\n",
    "        for i in range(25, 101):\n",
    "            tmp = random_df[random_df['len'] == i]\n",
    "            tmp.sort_values('total_reads', inplace=True, ascending=False)\n",
    "            tmp.reset_index(inplace=True, drop=True)\n",
    "            random_df_test = random_df_test.append(tmp.iloc[:100])\n",
    "\n",
    "        human_df_test = pd.DataFrame(columns=human_df.columns)\n",
    "        for i in range(25, 101):\n",
    "            tmp = human_df[human_df['len'] == i]\n",
    "            tmp.sort_values('total_reads', inplace=True, ascending=False)\n",
    "            tmp.reset_index(inplace=True, drop=True)\n",
    "            human_df_test = human_df_test.append(tmp.iloc[:100])\n",
    "\n",
    "        train_df = pd.concat([random_df, random_df_test]).drop_duplicates(keep=False)  # 去重后 76319\n",
    "        self.scaler = preprocessing.StandardScaler()\n",
    "        self.scaler.fit(train_df.loc[:, label_col].values.reshape(-1, 1))\n",
    "        train_df.loc[:,'scaled_rl'] = self.scaler.transform(train_df.loc[:, label_col].values.reshape(-1, 1))\n",
    "        random_df_test.loc[:, 'scaled_rl'] = self.scaler.transform(random_df_test.loc[:, label_col].values.reshape(-1, 1))\n",
    "        human_df_test.loc[:, 'scaled_rl'] = self.scaler.transform(human_df_test.loc[:, label_col].values.reshape(-1, 1))\n",
    "\n",
    "        if self.set_name == \"train\":\n",
    "            set_df = train_df\n",
    "        elif self.set_name == \"valid\":\n",
    "            set_df = random_df_test\n",
    "        else:\n",
    "            set_df = human_df_test\n",
    "\n",
    "        seq = set_df[seq_col].values\n",
    "        scaled_rl = set_df['scaled_rl'].values\n",
    "        selected_indices = set_df[\"ori_index\"].values\n",
    "\n",
    "        self.name = selected_indices\n",
    "\n",
    "        data_paths = {\"name\": self.name}\n",
    "        for itemname in data_items:\n",
    "            # input\n",
    "            if itemname == \"seq\":\n",
    "                data_paths[itemname] = seq\n",
    "            elif itemname == \"emb-rnafm\":\n",
    "                pass\n",
    "            # label\n",
    "            elif itemname == \"rl\":\n",
    "                data_paths[itemname] = scaled_rl\n",
    "            else:\n",
    "                raise Exception(\"Unknown data item name {}\".format(itemname))\n",
    "\n",
    "\n",
    "        print(\"{} Dataset Num Samples: {} \".format(self.set_name, set_df[\"len\"].shape[0]))\n",
    "        #print(\"Length-Frequency Table\")\n",
    "        #print(set_df[\"len\"].describe())  # value_counts())\n",
    "        return data_paths, set_df\n",
    "    \n",
    "    \n",
    "class BatchConverter(object):\n",
    "    \"\"\"\n",
    "    Callable to convert an unprocessed (labels + strings) batch to a processed (labels + tensor) batch.\n",
    "    \"\"\"\n",
    "    def __init__(self, alphabet, data_type=\"seq\",):\n",
    "        \"\"\"\n",
    "        :param alphabet:\n",
    "        :param data_type: seq, msa\n",
    "        \"\"\"\n",
    "        self.alphabet = alphabet\n",
    "        self.data_type = data_type.split(\"+\")\n",
    "\n",
    "    def __call__(self, raw_data, raw_anns=None):\n",
    "        \"\"\"\n",
    "        :param raw_data: each element in raw data should contain (description, seq)\n",
    "        :param raw_anns:\n",
    "        :return:\n",
    "        \"\"\"\n",
    "        # creat a new batch of data tensors\n",
    "        data = {}\n",
    "        for key in raw_data.keys():\n",
    "            if key == \"seq\":\n",
    "                labels, strs, tokens = self.__call_seq__(raw_data[\"seq\"])\n",
    "\n",
    "                data[\"description\"] = labels\n",
    "                data[\"string\"] = strs\n",
    "                data[\"token\"] = tokens\n",
    "                data[\"depth\"] = [1] * len(strs)\n",
    "                data[\"length\"] = [len(s) for s in strs]\n",
    "            else:\n",
    "                if isinstance(raw_data[key][0], str):\n",
    "                    data[key] = raw_data[key]\n",
    "                elif isinstance(raw_data[key][0], np.ndarray):\n",
    "                    try:   # same length\n",
    "                        data[key] = torch.Tensor(raw_data[key])\n",
    "                    except:\n",
    "                        # here we padding them with 0 for consistance with cnn's padding, which is different with ann's padding\n",
    "                        data[key] = torch.Tensor(self.__padding_numpy_matrix__(raw_data[key], data[\"length\"], pad_idx=0))\n",
    "                elif isinstance(raw_data[key][0], float) or isinstance(raw_data[key][0], int):\n",
    "                    data[key] = torch.Tensor(raw_data[key])\n",
    "\n",
    "\n",
    "        # creat a new batch of ann tensors\n",
    "        if raw_anns is not None:\n",
    "            anns = {}\n",
    "            for key in raw_anns.keys():\n",
    "                if isinstance(raw_anns[key][0], str):\n",
    "                    anns[key] = raw_anns[key]\n",
    "                elif isinstance(raw_anns[key][0], np.ndarray):\n",
    "                    try:   # same length\n",
    "                        anns[key] = torch.Tensor(raw_anns[key])\n",
    "                    except:\n",
    "                        anns[key] = torch.Tensor(self.__padding_numpy_matrix__(raw_anns[key], data[\"length\"]))\n",
    "                elif isinstance(raw_anns[key][0], float) or isinstance(raw_anns[key][0], int):\n",
    "                    anns[key] = torch.Tensor(raw_anns[key])\n",
    "        else:\n",
    "            anns = None\n",
    "\n",
    "        return data, anns\n",
    "\n",
    "    def __call_seq__(self, raw_batch: Sequence[Tuple[str, str]]):\n",
    "        # RoBERTa uses an eos token, while ESM-1 does not.\n",
    "        batch_size = len(raw_batch)\n",
    "        max_len = max(len(seq_str) for _, seq_str in raw_batch)\n",
    "        tokens = torch.empty(\n",
    "            (\n",
    "                batch_size,\n",
    "                max_len\n",
    "                + int(self.alphabet.prepend_bos)\n",
    "                + int(self.alphabet.append_eos),\n",
    "            ),\n",
    "            dtype=torch.int64,\n",
    "        )\n",
    "        tokens.fill_(self.alphabet.padding_idx)\n",
    "        labels = []\n",
    "        strs = []\n",
    "\n",
    "        for i, (label, seq_str) in enumerate(raw_batch):\n",
    "            labels.append(label)\n",
    "            strs.append(seq_str)\n",
    "            if self.alphabet.prepend_bos:\n",
    "                tokens[i, 0] = self.alphabet.cls_idx\n",
    "            seq = torch.tensor(\n",
    "                [self.alphabet.get_idx(s) for s in seq_str], dtype=torch.int64\n",
    "            )\n",
    "            tokens[\n",
    "            i,\n",
    "            int(self.alphabet.prepend_bos): len(seq_str)\n",
    "                                            + int(self.alphabet.prepend_bos),\n",
    "            ] = seq\n",
    "            if self.alphabet.append_eos:\n",
    "                tokens[\n",
    "                    i, len(seq_str) + int(self.alphabet.prepend_bos)\n",
    "                ] = self.alphabet.eos_idx\n",
    "\n",
    "        return labels, strs, tokens\n",
    "    \n",
    "def LofD_to_DofL(raw_batch):\n",
    "    \"\"\"\n",
    "    list of dict to dict of list\n",
    "    :param raw_batch:\n",
    "    :return:\n",
    "    \"\"\"\n",
    "    batch_size = len(raw_batch)\n",
    "    example = raw_batch[0]\n",
    "    new_batch = {}\n",
    "    for key in example.keys():\n",
    "        new_batch[key] = []\n",
    "        for i in range(batch_size):\n",
    "            new_batch[key].append(raw_batch[i][key])\n",
    "    return new_batch\n",
    "\n",
    "def build_collate_fn(alphabet, data_type):\n",
    "    batch_converter = BatchConverter(alphabet, data_type)\n",
    "    def collate_fn(batch):\n",
    "        if len(batch[0]) == 1:\n",
    "            data = zip(*batch)\n",
    "            data = LofD_to_DofL(data)\n",
    "            data, anns = batch_converter(data)\n",
    "            anns = None\n",
    "        elif len(batch[0]) == 2:\n",
    "            data, anns = zip(*batch)\n",
    "            data = LofD_to_DofL(data)\n",
    "            anns = LofD_to_DofL(anns)\n",
    "            data, anns = batch_converter(data, anns)\n",
    "        else:\n",
    "            raise Exception(\"Unexpected Num of Components in a Batch\")\n",
    "        return data, anns\n",
    "\n",
    "    return collate_fn\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "690cf64c",
   "metadata": {},
   "source": [
    "### (2) generate dataloaders"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "7029caaa",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/user/liyu/miniconda3/envs/jychen/lib/python3.8/site-packages/pandas/util/_decorators.py:311: SettingWithCopyWarning: \n",
      "A value is trying to be set on a copy of a slice from a DataFrame\n",
      "\n",
      "See the caveats in the documentation: https://pandas.pydata.org/pandas-docs/stable/user_guide/indexing.html#returning-a-view-versus-a-copy\n",
      "  return func(*args, **kwargs)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train Dataset Num Samples: 76319 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/user/liyu/miniconda3/envs/jychen/lib/python3.8/site-packages/pandas/util/_decorators.py:311: SettingWithCopyWarning: \n",
      "A value is trying to be set on a copy of a slice from a DataFrame\n",
      "\n",
      "See the caveats in the documentation: https://pandas.pydata.org/pandas-docs/stable/user_guide/indexing.html#returning-a-view-versus-a-copy\n",
      "  return func(*args, **kwargs)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "valid Dataset Num Samples: 7600 \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/user/liyu/miniconda3/envs/jychen/lib/python3.8/site-packages/pandas/util/_decorators.py:311: SettingWithCopyWarning: \n",
      "A value is trying to be set on a copy of a slice from a DataFrame\n",
      "\n",
      "See the caveats in the documentation: https://pandas.pydata.org/pandas-docs/stable/user_guide/indexing.html#returning-a-view-versus-a-copy\n",
      "  return func(*args, **kwargs)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "test Dataset Num Samples: 7600 \n"
     ]
    }
   ],
   "source": [
    "root_path = \"./\"\n",
    "train_set =  Human_5Prime_UTR_VarLength(root=root_path, set_name=\"train\", input_items=input_items, )\n",
    "val_set =  Human_5Prime_UTR_VarLength(root=root_path, set_name=\"valid\", input_items=input_items)\n",
    "test_set =  Human_5Prime_UTR_VarLength(root=root_path, set_name=\"test\", input_items=input_items)\n",
    "\n",
    "collate_fn = build_collate_fn(alphabet, train_set.data_type)\n",
    "\n",
    "num_workers = 4\n",
    "train_batch_size = 64\n",
    "train_loader = DataLoader(\n",
    "    train_set, batch_size=train_batch_size, #sampler=RandomSampler(train_set, replacement=False),\n",
    "    num_workers=num_workers, collate_fn=collate_fn, drop_last=False\n",
    ")\n",
    "\n",
    "val_batch_size = train_batch_size #* 4\n",
    "val_loader = DataLoader(\n",
    "    val_set, batch_size=val_batch_size, #sampler=RandomSampler(val_set, replacement=False),\n",
    "    num_workers=num_workers, collate_fn=collate_fn, drop_last=False\n",
    ")\n",
    "\n",
    "test_batch_size = train_batch_size #* 4\n",
    "test_loader = DataLoader(\n",
    "    test_set, batch_size=test_batch_size, #sampler=RandomSampler(test_set, replacement=False),\n",
    "    num_workers=num_workers, collate_fn=collate_fn, drop_last=False\n",
    ")\n",
    "\n",
    "scaler = train_set.scaler"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0b62c53c",
   "metadata": {},
   "source": [
    "## 3. Training Model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4469a0c7",
   "metadata": {},
   "source": [
    "### (1) define eval function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "d015847c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def model_eval(data_loader, i_epoch, set_name=\"unknown\"):\n",
    "    all_losses = []\n",
    "    #pbar = tqdm(data_loader, desc=\"Epoch {}, {} Set - real MSE: {}\".format(i_epoch,set_name, \"NaN\"), ncols=100)\n",
    "    #for index, (data, anns) in enumerate(pbar):\n",
    "    for index, (data, anns) in enumerate(data_loader):\n",
    "        backbone.eval()\n",
    "        utr_func_predictor.eval()\n",
    "        with torch.no_grad():\n",
    "            x = data[\"token\"].to(device)\n",
    "            true_rl = anns[\"rl\"].to(device)            \n",
    "            inputs = {}\n",
    "            results = {}\n",
    "            if \"seq\" in input_items:\n",
    "                inputs[\"seq\"] = x         \n",
    "            if \"emb-rnafm\" in input_items:\n",
    "                results = backbone(x, need_head_weights=False, repr_layers=[12], return_contacts=False)\n",
    "                inputs[\"emb-rnafm\"] = results[\"representations\"][12]    \n",
    "            results[\"rl\"] = utr_func_predictor(x, inputs)\n",
    "            #pds = scaler.inverse_transform(results[\"rl\"].detach().cpu().numpy())\n",
    "            #gts = scaler.inverse_transform(true_rl.detach().cpu().numpy())\n",
    "            #pds = torch.Tensor(pds)\n",
    "            #gts = torch.Tensor(gts)            \n",
    "        \n",
    "            #losses = criterion(pds, gts)  \n",
    "            losses = criterion(results[\"rl\"], true_rl)  \n",
    "            all_losses.append(losses.detach().cpu())     \n",
    "\n",
    "    avg_loss = torch.cat(all_losses, dim=0).mean()\n",
    "    print(\"Epoch {}, Evaluation on {} Set - MSE loss: {:.3f}\".format(i_epoch, set_name, avg_loss))\n",
    "    \n",
    "    return avg_loss"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b3561221",
   "metadata": {},
   "source": [
    "### (2) training process"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33bd1c35",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 1, Train Set - MSE loss: 0.531: 100%|█████████████████████| 1193/1193 [02:57<00:00,  6.72it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1, Evaluation on Random Set - MSE loss: 0.467\n",
      "--------- Model: CNN_seq_emb-rnafm, Best Epoch 1, Best MSE 0.467\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 2, Train Set - MSE loss: 0.455: 100%|█████████████████████| 1193/1193 [02:58<00:00,  6.70it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 2, Evaluation on Random Set - MSE loss: 0.495\n",
      "--------- Model: CNN_seq_emb-rnafm, Best Epoch 1, Best MSE 0.467\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 3, Train Set - MSE loss: 0.347:  78%|█████████████████     | 927/1193 [02:04<00:48,  5.48it/s]"
     ]
    }
   ],
   "source": [
    "n_epoches = 50\n",
    "best_mse = 10\n",
    "best_epoch = 0\n",
    "\n",
    "for i_e in range(1, n_epoches+1):\n",
    "    all_losses = []\n",
    "    n_sample = 0\n",
    "    n_iter = len(train_loader)\n",
    "\n",
    "    pbar = tqdm(train_loader, desc=\"Epoch {}, Train Set - MSE loss: {}\".format(i_e, \"NaN\"), ncols=100)\n",
    "    for index, (data, anns) in enumerate(pbar):\n",
    "        backbone.eval()\n",
    "        utr_func_predictor.train()\n",
    "        x = data[\"token\"].to(device)\n",
    "        true_rl = anns[\"rl\"].to(device)        \n",
    "        \n",
    "        inputs = {}\n",
    "        results = {}\n",
    "        if \"seq\" in input_items:\n",
    "            inputs[\"seq\"] = x         \n",
    "        if \"emb-rnafm\" in  input_items:            \n",
    "            with torch.no_grad():\n",
    "                results = backbone(x, need_head_weights=False, repr_layers=[12], return_contacts=False)            \n",
    "            inputs[\"emb-rnafm\"] = results[\"representations\"][12]                \n",
    "        results[\"rl\"] = utr_func_predictor(x, inputs)\n",
    "        losses = criterion(results[\"rl\"], true_rl)\n",
    "        batch_loss = losses.mean()\n",
    "        batch_loss.backward()\n",
    "        optimizer.step()\n",
    "        optimizer.zero_grad()\n",
    "    \n",
    "        all_losses.append(losses.detach().cpu())\n",
    "        current_avg_loss = torch.cat(all_losses, dim=0).mean()\n",
    "        \n",
    "        pbar.set_description(\"Epoch {}, Train Set - MSE loss: {:.3f}\".format(i_e, current_avg_loss))\n",
    "        #if index % 1 == 0:\n",
    "        #    print(\"epoch: {}, iter:{}/{} MSE loss: {}\".format(i_e, index+1, iters, all_loss/n_sample ))\n",
    "    \n",
    "        #if index+1 > 10:\n",
    "        #    break    \n",
    "    \n",
    "    random_mse = model_eval(val_loader, i_e, set_name=\"Random\")\n",
    "    \n",
    "    if random_mse < best_mse:\n",
    "        best_epoch = i_e\n",
    "        best_mse = random_mse\n",
    "        torch.save(utr_func_predictor.state_dict(), \"result/{}_best_utr_predictor.pth\".format(model_name))\n",
    "    print(\"--------- Model: {}, Best Epoch {}, Best MSE {:.3f}\".format(model_name, best_epoch, best_mse))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d920d2ff",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
